package top.lingkang.finalgateway.monitor;

import cn.hutool.core.net.Ipv4Util;
import cn.hutool.core.util.StrUtil;
import org.noear.solon.Solon;
import top.lingkang.finalgateway.admin.entity.IpListEntity;
import top.lingkang.finalgateway.admin.mapper.IpListLogMapper;
import top.lingkang.finalgateway.admin.mapper.IpListMapper;
import top.lingkang.finalgateway.admin.vo.RateIpVO;
import top.lingkang.finalgateway.utils.CommonUtils;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;

/**
 * @author lingkang
 * Created by 2024/5/17
 */
public class IpListCheck {
    private final int port;
    private final IpListMapper ipListMapper;
    private final IpListLogMapper ipListLogMapper;
    private final HashSet<String> allow = new HashSet<>(512);
    private final HashSet<String> forbidden = new HashSet<>(512);
    private volatile List<String> white = new ArrayList<>();
    private List<String> black = new ArrayList<>();

    public IpListCheck(int port) {
        this.port = port;
        ipListMapper = Solon.context().getBean(IpListMapper.class);
        loadIpList();
        ipListLogMapper = Solon.context().getBean(IpListLogMapper.class);
    }

    public void loadIpList() {
        IpListEntity entity = ipListMapper.selectById(port);
        if (entity != null) {
            white = CommonUtils.toIpList(entity.getWhite());
            black = CommonUtils.toIpList(entity.getBlack());
        } else {
            white = new ArrayList<>();
            black = white;
        }
    }

    public void clearCache() {
        white.clear();
        black.clear();
        allow.clear();
        forbidden.clear();

        // 清理完后重新加载
        loadIpList();
    }

    public boolean checkIp(String ip) {
        if (allow.contains(ip))
            return true;
        if (forbidden.contains(ip))
            return false;

        // 检测白名单
        if (!white.isEmpty()) {
            for (String reg : white)
                if (Ipv4Util.matches(reg, ip)) {
                    allow.add(ip);
                    return true;
                }
            return false;
        }

        // 黑名单
        for (String reg : black)
            if (Ipv4Util.matches(reg, ip)) {
                forbidden.add(ip);
                return false;
            }

        // 允许
        allow.add(ip);
        return true;
    }

    /**
     * 添加黑名单
     */
    public synchronized void addBlack(RateIpVO vo) {
        if (!isWhiteIp(vo.getIp())) {
            allow.remove(vo.getIp());
        } else {
            allow.add(vo.getIp());
            return;
        }
        // 已经在黑名单了
        if (black.contains(vo.getIp()))
            return;

        forbidden.add(vo.getIp());
        IpListEntity entity = ipListMapper.selectById(port);
        if (entity == null) {
            entity = new IpListEntity();
            entity.setPort(port);
            entity.setBlack(vo.getIp());
            ipListMapper.insert(entity);
            return;
        }
        if (StrUtil.isNotBlank(entity.getBlack()))
            entity.setBlack(entity.getBlack() + "\n" + vo.getIp());
        else
            entity.setBlack(vo.getIp());
        ipListMapper.updateById(entity);
        ipListLogMapper.add(port, "安全防御", "ip: " + vo.getIp() + " 频繁访问端口, 2分钟内次数：" + vo.getNumber());
    }

    public boolean isWhiteIp(String ip) {
        for (String reg : white)
            if (Ipv4Util.matches(reg, ip)) {
                allow.add(ip);
                return true;
            }
        return false;
    }
}
